/*
 * MythMobile, a Django interface to MythTV for mobile devices
 * Copyright (C) 2011 Andrew Wilkinson (andrewjwilkinson@gmail.com)
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Portions of this file are Copyright (c) 2009 Chase Douglas.
 * These portions are taken from segmenter, available at http://svn.assembla.com/svn/legend/segmenter/segmenter.c
 * which were released under the GNU GPL.
 */

#include <time.h>
#include <Python.h>

#include "libavformat/avformat.h"

#define BUF_SIZE (4096*2)

static AVStream *add_output_stream(AVFormatContext *output_format_context, AVStream *input_stream) {
    AVCodecContext *input_codec_context;
    AVCodecContext *output_codec_context;
    AVStream *output_stream;

    output_stream = av_new_stream(output_format_context, 0);
    if (!output_stream) {
        fprintf(stderr, "Could not allocate stream\n");
        exit(1);
    }

    input_codec_context = input_stream->codec;
    output_codec_context = output_stream->codec;

    output_codec_context->codec_id = input_codec_context->codec_id;
    output_codec_context->codec_type = input_codec_context->codec_type;
    output_codec_context->codec_tag = input_codec_context->codec_tag;
    output_codec_context->bit_rate = input_codec_context->bit_rate;
    output_codec_context->extradata = input_codec_context->extradata;
    output_codec_context->extradata_size = input_codec_context->extradata_size;

    if(av_q2d(input_codec_context->time_base) * input_codec_context->ticks_per_frame > av_q2d(input_stream->time_base) && av_q2d(input_stream->time_base) < 1.0/1000) {
        output_codec_context->time_base = input_codec_context->time_base;
        output_codec_context->time_base.num *= input_codec_context->ticks_per_frame;
    }
    else {
        output_codec_context->time_base = input_stream->time_base;
    }

    switch (input_codec_context->codec_type) {
        case AVMEDIA_TYPE_AUDIO:
            output_codec_context->channel_layout = input_codec_context->channel_layout;
            output_codec_context->sample_rate = input_codec_context->sample_rate;
            output_codec_context->channels = input_codec_context->channels;
            output_codec_context->frame_size = input_codec_context->frame_size;
            if ((input_codec_context->block_align == 1 && input_codec_context->codec_id == CODEC_ID_MP3) || input_codec_context->codec_id == CODEC_ID_AC3) {
                output_codec_context->block_align = 0;
            }
            else {
                output_codec_context->block_align = input_codec_context->block_align;
            }
            break;
        case AVMEDIA_TYPE_VIDEO:
            output_codec_context->pix_fmt = input_codec_context->pix_fmt;
            output_codec_context->width = input_codec_context->width;
            output_codec_context->height = input_codec_context->height;
            output_codec_context->has_b_frames = input_codec_context->has_b_frames;

            if (output_format_context->oformat->flags & AVFMT_GLOBALHEADER) {
                output_codec_context->flags |= CODEC_FLAG_GLOBAL_HEADER;
            }
            break;
    default:
        break;
    }

    return output_stream;
}

typedef struct {
    PyObject_HEAD

    AVInputFormat* ifmt;
    AVFormatContext* ic;
    AVOutputFormat* ofmt;
    AVFormatContext* oc;
    ByteIOContext* bioc;
    char* buffer;
    int audio_index;
    int video_index;
    AVStream *video_st;
    AVStream *audio_st;
    AVCodec *codec;
    PyObject* processed_buffer;
    int packet;
    int packet_unread;
    PyObject* new_filename_func;
    int max_segment_time;
    double segment_time;
    double prev_segment_time;
} mythmobile_segmentObject;

int mythmobile_segment_stream_read(void *opaque, uint8_t *buf, int buf_size) {
    mythmobile_segmentObject* self = (mythmobile_segmentObject*)opaque;

    PyObject* packet = PyInt_FromLong(self->packet);

    int part_cpy = 0;
    int cpy_size = 0;
    int delay = 0;
    while(cpy_size == 0 && delay < 10) {
        if(PyDict_Contains(self->processed_buffer, packet) > 0) {
            PyObject* data = PyDict_GetItem(self->processed_buffer, packet);

            int data_size = 0;
            char* raw_data = NULL;

            PyString_AsStringAndSize(data, &raw_data, &data_size);

            cpy_size = data_size - self->packet_unread;
            if(cpy_size > buf_size) {
                part_cpy = 1;
                cpy_size = buf_size;
            }
            memcpy(buf, raw_data + self->packet_unread, cpy_size);

            if(part_cpy) {
                self->packet_unread += cpy_size;
            } else {
                PyDict_DelItem(self->processed_buffer, packet);
                self->packet++;
                self->packet_unread = 0;
            }
        } else {
            Py_BEGIN_ALLOW_THREADS
            sleep(1);
            delay += 1;
            Py_END_ALLOW_THREADS
        }
    }

    Py_DECREF(packet);

    return cpy_size;
}

int mythmobile_segment_stream_write(void *opaque, uint8_t *buf, int buf_size) {
    return 0;
}

int64_t mythmobile_segment_stream_seek(void *opaque, int64_t offset, int whence) {
    return -1;
}

static PyObject *
mythmobile_segmentObject_new(PyTypeObject *type, PyObject *args, PyObject *kwds) {
    mythmobile_segmentObject* self;

    self = (mythmobile_segmentObject*)type->tp_alloc(type, 0);
    if (self != NULL) {
        self->ifmt = NULL;
        self->ic = NULL;
        self->ofmt = NULL;
        self->oc = NULL;
        self->bioc = NULL;
        self->buffer = NULL;
        self->audio_index = -1;
        self->video_index = -1;
        self->video_st = NULL;
        self->audio_st = NULL;
        self->codec = NULL;
        self->processed_buffer = NULL;
        self->packet = 0;
        self->packet_unread = 0;
        self->new_filename_func = NULL;
        self->max_segment_time = 0;
        self->segment_time = 0;
        self->prev_segment_time = 0;
    }

    return (PyObject*)self;
}

static int
mythmobile_segmentObject_init(mythmobile_segmentObject* self, PyObject *args, PyObject *kwds) {
    char* input_filename;
    int i;

    if (!PyArg_ParseTuple(args, "sOOi", &input_filename, &self->new_filename_func, &self->processed_buffer, &self->max_segment_time))
        return -1;

    Py_INCREF(self->new_filename_func);
    Py_INCREF(self->processed_buffer);

    self->ifmt = av_find_input_format("mpegts"); 
    if(!self->ifmt) {
        PyErr_SetString(PyExc_SystemError, "Failed to find input format 'mpegts'.");
        return -1;
    }

    self->buffer = (char*)malloc(BUF_SIZE + FF_INPUT_BUFFER_PADDING_SIZE);
    self->bioc = malloc(sizeof(ByteIOContext));
    self->bioc->is_streamed = 1;

    if(init_put_byte(self->bioc, self->buffer, BUF_SIZE, 0, self, &mythmobile_segment_stream_read, &mythmobile_segment_stream_write, &mythmobile_segment_stream_seek)) {
        PyErr_SetString(PyExc_SystemError, "Failed to initalize ByteIOContext.");
        return -1;
    }

    if(av_open_input_stream(&self->ic, self->bioc, input_filename, self->ifmt, NULL) != 0) {
        PyErr_SetString(PyExc_ValueError, "Failed to open input file. Is it an mpegts file?");
        return -1;
    }

    if(av_find_stream_info(self->ic) < 0) {
        PyErr_SetString(PyExc_ValueError, "Failed to get stream info.");
        return -1;
    }

    self->ofmt = av_guess_format("mpegts", NULL, NULL);
    if(!self->ofmt) {
        PyErr_SetString(PyExc_SystemError, "Could not find MPEG-TS muxer.");
        return -1;
    }

    self->oc = avformat_alloc_context();
    if(!self->oc) {
        PyErr_SetString(PyExc_SystemError, "Could not allocated output context.");
        return -1;
    }
    self->oc->oformat = self->ofmt;

    for(i = 0; i < self->ic->nb_streams && (self->video_index < 0 || self->audio_index < 0); i++) {
        switch (self->ic->streams[i]->codec->codec_type) {
            case AVMEDIA_TYPE_VIDEO:
                self->video_index = i;
                self->ic->streams[i]->discard = AVDISCARD_NONE;
                self->video_st = add_output_stream(self->oc, self->ic->streams[i]);
                break;
            case AVMEDIA_TYPE_AUDIO:
                self->audio_index = i;
                self->ic->streams[i]->discard = AVDISCARD_NONE;
                self->audio_st = add_output_stream(self->oc, self->ic->streams[i]);
                break;
            default:
                self->ic->streams[i]->discard = AVDISCARD_ALL;
                break;
        }
    }

    if(av_set_parameters(self->oc, NULL) < 0) {
        PyErr_SetString(PyExc_SystemError, "Invalid output format parameters.");
        return -1;
    }

    self->codec = avcodec_find_decoder(self->video_st->codec->codec_id);
    if(!self->codec) {
        PyErr_SetString(PyExc_SystemError, "Could not find video decoder.");
        return -1;
    }

    if(avcodec_open(self->video_st->codec, self->codec) < 0) {
        PyErr_SetString(PyExc_SystemError, "Could not open video decoder.");
        return -1;
    }

    return 0;
}

static void
mythmobile_segmentObject_dealloc(mythmobile_segmentObject* self)
{
    Py_DECREF(self->new_filename_func);
    Py_DECREF(self->processed_buffer);

    if(self->ifmt != NULL) {
    }

    if(self->ic != NULL) {
    }

    if(self->ofmt != NULL) {
    }

    if(self->oc != NULL) {
    }

    if(self->bioc != NULL) {
    }

    if(self->buffer != NULL) {
        free(self->buffer);
    }

    self->ob_type->tp_free((PyObject*)self);
}

static PyObject* mythmobile_segmentObject_process(mythmobile_segmentObject* self) {
    AVPacket packet;
    int decode_done;
    int write_ret;
    int new_file = 1;

    while(1) {
        if(new_file) {
            PyObject* new_filename = PyIter_Next(self->new_filename_func);
            if(new_filename == NULL) {
                if(!PyErr_Occurred()) {
                    PyErr_SetString(PyExc_TypeError, "Filename iterator ended.");
                }

                return NULL;
            }

            if (url_fopen(&self->oc->pb, PyString_AsString(new_filename), URL_WRONLY) < 0) {
                PyErr_SetString(PyExc_SystemError, "Could not open new file.");
                break;
            }

            if(av_write_header(self->oc)) {
                PyErr_SetString(PyExc_SystemError, "Could not write mpegts header to output file.");
                return NULL;
            }

            self->prev_segment_time = self->segment_time;
            self->segment_time = 0;
            new_file = 0;
            
            Py_DECREF(new_filename);
        }

        decode_done = av_read_frame(self->ic, &packet);
        if (decode_done < 0) {
            PyErr_SetString(PyExc_SystemError, "Couldn't decode frame.");
            return NULL;
        }

        if (av_dup_packet(&packet) < 0) {
            av_free_packet(&packet);
            PyErr_SetString(PyExc_SystemError, "Couldn't duplicate packet.");
            return NULL;
        }

        if (packet.stream_index == self->video_index && (packet.flags & AV_PKT_FLAG_KEY)) {
            self->segment_time = (double)self->video_st->pts.val * self->video_st->time_base.num / self->video_st->time_base.den;
        }
        else if (self->video_index < 0) {
            self->segment_time = (double)self->audio_st->pts.val * self->audio_st->time_base.num / self->audio_st->time_base.den;
        }

        write_ret = av_interleaved_write_frame(self->oc, &packet);
        if (write_ret < 0) {
            fprintf(stderr, "Warning: Could not write frame of stream %i\n", write_ret);
        } else if (write_ret > 0) {
            PyErr_SetString(PyExc_SystemError, "End of stream requested.");
            av_free_packet(&packet);
            return NULL;
        }

        if((self->segment_time - self->prev_segment_time) > self->max_segment_time) {
            put_flush_packet(self->oc->pb);
            url_fclose(self->oc->pb);

            new_file = 1;
        }

        //self->prev_segment_time = self->segment_time;
    }

    Py_INCREF(Py_None);
    return Py_None;
}

static PyMethodDef mythmobile_segmentObject_methods[] = {
    {"process", (PyCFunction)mythmobile_segmentObject_process, METH_NOARGS,
     "Process some video data."
    },
    {NULL}  /* Sentinel */
};

static PyTypeObject mythmobile_segmentType = {
    PyObject_HEAD_INIT(NULL)
    0,                         /*ob_size*/
    "_mythmobile.segment",     /*tp_name*/
    sizeof(mythmobile_segmentObject), /*tp_basicsize*/
    0,                         /*tp_itemsize*/
    (destructor)&mythmobile_segmentObject_dealloc, /*tp_dealloc*/
    0,                         /*tp_print*/
    0,                         /*tp_getattr*/
    0,                         /*tp_setattr*/
    0,                         /*tp_compare*/
    0,                         /*tp_repr*/
    0,                         /*tp_as_number*/
    0,                         /*tp_as_sequence*/
    0,                         /*tp_as_mapping*/
    0,                         /*tp_hash */
    0,                         /*tp_call*/
    0,                         /*tp_str*/
    0,                         /*tp_getattro*/
    0,                         /*tp_setattro*/
    0,                         /*tp_as_buffer*/
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,        /*tp_flags*/
    "MythMobile Segment",     /* tp_doc */
    0,                     /* tp_traverse */
    0,                     /* tp_clear */
    0,                     /* tp_richcompare */
    0,                     /* tp_weaklistoffset */
    0,                     /* tp_iter */
    0,                     /* tp_iternext */
    mythmobile_segmentObject_methods,             /* tp_methods */
    0,             /* tp_members */
    0,                         /* tp_getset */
    0,                         /* tp_base */
    0,                         /* tp_dict */
    0,                         /* tp_descr_get */
    0,                         /* tp_descr_set */
    0,                         /* tp_dictoffset */
    (initproc)mythmobile_segmentObject_init,      /* tp_init */
    0,                         /* tp_alloc */
    mythmobile_segmentObject_new,                 /* tp_new */
};

static PyMethodDef MythMobileMethods[] = {
    {NULL, NULL, 0, NULL}        /* Sentinel */
};

PyMODINIT_FUNC
init_mythmobile(void)
{
    PyObject *m;

    av_register_all();

    if (PyType_Ready(&mythmobile_segmentType) < 0)
        return;

    m = Py_InitModule("_mythmobile", MythMobileMethods);
    if (m == NULL)
        return;

    Py_INCREF(&mythmobile_segmentType);
    PyModule_AddObject(m, "Segment", (PyObject *)&mythmobile_segmentType);
}
